package com.wfm.leetcode.editor.cn;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 输入
 *  {1, 3, 5, 7, 9},
 *  {2, 4, 6, 8, 10},
 *  {3, 4, 5, 6, 7}
 *  输出
 *  [10, 9, 8, 7, 7]
 */
public class P二维数组的TopK {
    public static void main(String[] args) {
        int[][] matrix = {
                {1, 3, 5, 7, 9},
                {2, 4, 6, 8, 10},
                {3, 4, 5, 6, 7}
        };
        int k = 5;

        List<Integer> topK = findTopK(matrix, k);
        System.out.println(topK); // Output: [10, 9, 8, 7, 7]
    }
    public static List<Integer> findTopK(int[][] matrix, int k) {
        List<Integer> result = new ArrayList<>();
        if (matrix == null || matrix.length == 0 || k <= 0) {
            return result;
        }

        PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> b[0] - a[0]);

        // Initialize the max heap with the last element from each row
        for (int i = 0; i < matrix.length; i++) {
            int col = matrix[i].length;
            if (col > 0) {
                maxHeap.offer(new int[]{matrix[i][col - 1], i, col - 1});
            }
        }

        while (!maxHeap.isEmpty() && k > 0) {
            int[] top = maxHeap.poll();
            int val = top[0];
            int row = top[1];
            int col = top[2];

            result.add(val);
            k--;

            // Move to the previous element in the same row if available
            if (col - 1 >= 0) {
                maxHeap.offer(new int[]{matrix[row][col - 1], row, col - 1});
            }
        }

        return result;
    }
}
